#!/usr/bin/env python

from __future__ import division

import functools, math, sys
import matplotlib.pyplot as plt


def choose_step(method, eps):
    h = 1.0
    x_vals_1, y_vals_1 = method(h)
    while h:
        h /= 2
        x_vals_2, y_vals_2 = method(h)
        shorten_y_vals = [y for i, y in enumerate(y_vals_2) if i % 2 == 0]
        dy_vals = [math.fabs(y1 - y2) for y1, y2 in zip(y_vals_1, shorten_y_vals)]
        dy_vals = [True if dy < eps else False for dy in dy_vals]
        if all(dy_vals):
            return h
        if any(dy_vals[1:]):
            print 'method converged with h = {0} at the {1} point'.format(h, dy_vals[1:].index(True) + 1)
            print 'disconverged at the {0} point'.format(dy_vals.index(False))
        x_vals_1, y_vals_1 = x_vals_2, y_vals_2
    # if we could not find step good enough
    return sys.float_info.min


def euler_method(f, x0, y0, length, h):
    x = float(x0)
    y = float(y0)
    x_vals = [x]
    y_vals = [y]

    while x < length - x0:
        y = y + h * f(x, y)
        x += h
        x_vals.append(x)
        y_vals.append(y)

    return x_vals, y_vals


def modified_euler_method(f, x0, y0, length, h):
    x = float(x0)
    y = float(y0)
    x_vals = [x]
    y_vals = [y]

    while x < length - x0:
        y = y + h * f(x + h / 2, y + h / 2 * f(x, y))
        #predicted_y = y + h * f(x, y)
        #y = y + h / 2 * (f(x, y) + f(x, predicted_y))
        x += h
        x_vals.append(x)
        y_vals.append(y)

    return x_vals, y_vals


def runge_kutta_method(f, x0, y0, length, h):
    x = float(x0)
    y = float(y0)
    x_vals = [x]
    y_vals = [y]

    while x < length - x0:
        k1 = h * f(x, y)
        k2 = h * f(x + h / 2, y + k1 / 2)
        k3 = h * f(x + h / 2, y + k2 / 2)
        k4 = h * f(x + h, y + k3)
        y = y + (k1 + 2 * k2 + 2 * k3 + k4) / 6
        x += h
        x_vals.append(x)
        y_vals.append(y)

    return x_vals, y_vals


def main():
    f = lambda x, y: (0.7 * (1 - math.pow(y, 2))) / ((1 + 1.5) * math.pow(x, 2) + math.pow(y, 2) + 1)
    x0 = 0
    y0 = 0
    length = 1
    eps = 0.001

    print 'Euler method:'
    h = choose_step(functools.partial(euler_method, f, x0, y0, length), eps)
    print 'step:', h
    x_vals, y_vals = euler_method(f, x0, y0, length, h)
    p1, = plt.plot(x_vals, y_vals)

    print 'Modified Euler method:'
    h = choose_step(functools.partial(modified_euler_method, f, x0, y0, length), eps)
    print 'step:', h
    x_vals, y_vals = modified_euler_method(f, x0, y0, length, h)
    p2, = plt.plot(x_vals, y_vals)

    print 'Runge-Kutta method:'
    h = choose_step(functools.partial(runge_kutta_method, f, x0, y0, length), eps)
    x_vals, y_vals = runge_kutta_method(f, x0, y0, length, h)
    print 'step:', h
    p3, = plt.plot(x_vals, y_vals)

    plt.legend( (p1, p2, p3), ('Euler', 'Modified Euler', 'Runge-Kutta') )
    plt.show()


if __name__ == '__main__':
    main()
